#include "tcp_config.hh"
#include "tcp_expectation.hh"
#include "tcp_fsm_test_harness.hh"
#include "tcp_header.hh"
#include "tcp_segment.hh"
#include "util.hh"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>

using namespace std;
using State = TCPTestHarness::State;

static constexpr unsigned NREPS = 32;

int main() {
  try {
    TCPConfig cfg{};
    cfg.recv_capacity = 65000;
    auto rd = get_random_generator();

    // loop segments back in a different order
    for (unsigned rep_no = 0; rep_no < NREPS; ++rep_no) {
      const WrappingInt32 rx_offset(rd());
      TCPTestHarness test_2 = TCPTestHarness::in_established(cfg, rx_offset - 1, rx_offset - 1);
      test_2.send_ack(rx_offset, rx_offset, 65000);

      string d(cfg.recv_capacity, 0);
      generate(d.begin(), d.end(), [&] { return rd(); });

      vector<TCPSegment> segs;
      size_t sendoff = 0;
      while (sendoff < d.size()) {
        const size_t len = min(d.size() - sendoff, static_cast<size_t>(rd()) % 8192);
        if (len == 0) {
          continue;
        }
        test_2.execute(Write{d.substr(sendoff, len)});
        test_2.execute(Tick(1));

        test_2.execute(ExpectSegmentAvailable{}, "test 2 failed: cannot read after write()");
        while (test_2.can_read()) {
          segs.emplace_back(test_2.expect_seg(ExpectSegment{}));
        }
        sendoff += len;
      }

      // resend them in shuffled order
      vector<size_t> seg_idx(segs.size(), 0);
      iota(seg_idx.begin(), seg_idx.end(), 0);
      shuffle(seg_idx.begin(), seg_idx.end(), rd);
      vector<TCPSegment> acks;
      for (auto idx : seg_idx) {
        test_2.execute(SendSegment{std::move(segs[idx])});
        test_2.execute(Tick(1));
        TCPSegment s = test_2.expect_seg(ExpectOneSegment{}.with_ack(true), "test 2 failed: no ACK after rcvd");
        acks.emplace_back(std::move(s));
        test_2.execute(ExpectNoSegment{}, "test 2 failed: double ACK?");
      }

      // send just the final ack
      test_2.execute(SendSegment{std::move(acks.back())});
      test_2.execute(ExpectNoSegment{}, "test 2 failed: ACK for ACK?");

      test_2.execute(ExpectData{}.with_data(d), "test 2 failed: wrong data after loopback");
    }
  } catch (const exception &e) {
    cerr << e.what() << endl;
    return EXIT_FAILURE;
  }

  return EXIT_SUCCESS;
}
